import os
import numpy as np
import re
from threading import Thread
from filelock import FileLock

# 预先计算的权重矩阵加载
weights = {
    'left_upper': np.load('E:\\wafer52\\weight\\Top_Left.npy'),
    'right_upper': np.load('E:\\wafer52\\weight\\Top_Right.npy'),
    'left_lower': np.load('E:\\wafer52\\weight\\Bottom_Left.npy'),
    'right_lower': np.load('E:\\wafer52\\weight\\Bottom_Right.npy')
}

def get_sub_flow_files(row, col):
    return [
        (row, col, 'left_upper'),
        (row, col + 1, 'right_upper'),
        (row + 1, col, 'left_lower'),
        (row + 1, col + 1, 'right_lower')
    ]

def process_flow_file(flow_file, flow_dir, output_dir, tile_size=256, full_size=512):
    match = re.search(r'tile_\((\d+)_(\d+),(\d+)_(\d+)\)\.npy', flow_file)
    if match:
        row1, col1, row2, col2 = map(int, match.groups())
        parent_path = os.path.join(flow_dir, flow_file)
        flow = np.load(parent_path)

        sub_flows = get_sub_flow_files(row1, col1)
        for (row, col, weight_key) in sub_flows:
            weight_matrix = weights[weight_key][:, :, np.newaxis]  # 扩展维度以匹配光流场
            if weight_key == 'left_upper':
                sub_flow = flow[:tile_size, :tile_size] * weight_matrix
            elif weight_key == 'right_upper':
                sub_flow = flow[:tile_size, tile_size:] * weight_matrix
            elif weight_key == 'left_lower':
                sub_flow = flow[tile_size:, :tile_size] * weight_matrix
            elif weight_key == 'right_lower':
                sub_flow = flow[tile_size:, tile_size:] * weight_matrix

            output_file = os.path.join(output_dir, f"tile_({row}_{col}).npy")
            lock_file = output_file + ".lock"
            with FileLock(lock_file):
                if os.path.exists(output_file):
                    existing_flow = np.load(output_file)
                    sub_flow += existing_flow
                np.save(output_file, sub_flow)
            print(f"Processed {output_file}")

def worker(flow_files, flow_dir, output_dir, tile_size, full_size):
    for flow_file in flow_files:
        process_flow_file(flow_file, flow_dir, output_dir, tile_size, full_size)

def merge_flows(flow_dir, output_dir, tile_size=256, full_size=512):
    os.makedirs(output_dir, exist_ok=True)
    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]

    num_threads = 8
    chunk_size = len(flow_files) // num_threads
    threads = []

    for i in range(num_threads):
        start_index = i * chunk_size
        end_index = (i + 1) * chunk_size if i < num_threads - 1 else len(flow_files)
        thread_files = flow_files[start_index:end_index]
        thread = Thread(target=worker, args=(thread_files, flow_dir, output_dir, tile_size, full_size))
        threads.append(thread)
        thread.start()

    for thread in threads:
        thread.join()

# 示例用法
flow_dir = 'E:\\wafer52\\11867_32nm_fine_align_512_dvf(1)'  # 替换为实际输入目录
output_dir = 'C:\\11867_32nm_fine_align_256_dvf_repair(3)'  # 替换为实际输出目录

merge_flows(flow_dir, output_dir, tile_size=256, full_size=512)

